from __future__ import with_statement

import bisect
import optparse
import sys
import math
import re

import numpy

import matplotlib
matplotlib.use("PS")
import matplotlib.figure
import matplotlib.backends.backend_ps
import pylab

import common
import wlc
import refolding_common


def m_to_nm(a):
    if numpy.isscalar(a):
        return a * 10 ** 9
    else:
        return (numpy.asarray(a) * 10**9).tolist()

def n_to_pn(a):
    if numpy.isscalar(a):
        return a * 10 ** 12
    else:
        return (numpy.asarray(a) * 10 ** 12).tolist()

class RefoldingTrace(object):
    def __init__(
        self,
        kb_t_over_p, k_n_per_m, hover_time_s, hover_z_m, hover_f_n, fetch_z_m, fetch_f_n,
        refolding_z_m, refolding_f_n, cl_m):

        common.check_equal(len(fetch_z_m), len(fetch_f_n))
        common.check_equal(len(refolding_z_m), len(refolding_f_n))
        self.kb_t_over_p = kb_t_over_p
        common.check_gt(k_n_per_m, 0)
        self.k_n_per_m = k_n_per_m
        self.hover_time_s = hover_time_s
        self.hover_z_m = hover_z_m
        self.hover_f_n = hover_f_n
        self.fetch_z_m = fetch_z_m
        self.fetch_f_n = fetch_f_n
        self.refolding_z_m = refolding_z_m
        self.refolding_f_n = refolding_f_n
        self.cl_m = cl_m

def _create_figure():
    y_factor=0.5
    bottom=0.19
    spp = matplotlib.figure.SubplotParams(left=0.18, bottom=bottom, right=0.955, wspace=0, hspace=0)
    _factor = 1.5
    figure = matplotlib.figure.Figure(figsize=(0.675 * 3.375 * _factor, 2.7 * _factor * y_factor), subplotpars=spp)
    return figure

def refolding_sample_curve(line, wlc_high_n, ignored_peaks_at_end, ax):
    trace = eval(line)
    ax.plot(m_to_nm(trace.fetch_z_m), n_to_pn(-numpy.array(trace.fetch_f_n)), color="red")
    high = wlc_high_n * 1e12
    amin = numpy.argmin(trace.refolding_f_n)
    start = amin - 100
    offset = high + 10
    ax.plot(
        m_to_nm(trace.refolding_z_m[start:]),
        map(lambda x : x + offset, n_to_pn(-numpy.array(trace.refolding_f_n[start:]))),
        color="black")
    cl_m = trace.cl_m[0:len(trace.cl_m) - ignored_peaks_at_end]
    last_color = "crimson"
    for (i, cl) in enumerate(cl_m + [cl_m[-1] + refolding_common.gb1_loop_length]):
        # Number of points will be equal to number of Angstroms.
        x = numpy.arange(0, cl * 0.99, 1e-10)  # numpy.array(sorted(set([z_m for z_m in trace.fetch_z_m if z_m < cl])), dtype=float)
        f = wlc.wlc(trace.kb_t_over_p, x / cl)
        z = x + f / trace.k_n_per_m
        n = bisect.bisect_right(f, high * 1e-12)
        ax.plot(m_to_nm(z[:n]), n_to_pn(f[:n]), color=last_color if i == len(cl_m) else "blue")
        ax.plot(
            m_to_nm(z[:n]), map(lambda x : x + offset, n_to_pn(f[:n])),
            color=last_color if i == len(cl_m) else "blue")
        if i == len(cl_m):
            ax.text(
                m_to_nm(z[n - 1]), n_to_pn(f[n - 1]) / 2.0, "%.1f nm" % (cl * 1e9),
                color=last_color, horizontalalignment="left", verticalalignment="center",
                family="monospace", size="x-small")
    total_modules = len(cl_m) + int((trace.refolding_z_m[amin] - cl_m[-1]) // refolding_common.gb1_loop_length)
    for (peak, cl) in enumerate(cl_m + [cl_m[-1] + refolding_common.gb1_loop_length * (1 + i) for i in range(total_modules - len(cl_m))]):
        x, y = cl * 0.87 * 1e9, offset + high + 50
        t = ax.text(x, y, repr(peak + 1), family="monospace", size="x-small")
        #ax.add_patch(matplotlib.patches.Circle((x, y), radius=50))
    vertical_segment_lower_limit = -max(trace.fetch_f_n) * 1e12
    vertical_segment_upper_limit = -trace.refolding_f_n[amin] * 1e12 + offset
    hover_color = "indigo"
    ax.plot(
        m_to_nm([trace.hover_z_m, trace.hover_z_m]),
        [vertical_segment_lower_limit, vertical_segment_upper_limit],
        linestyle="dashed", color=hover_color)
    ax.text(
        m_to_nm(trace.hover_z_m), vertical_segment_upper_limit + 150,
        "%.1f nm, %d ms" % (m_to_nm(trace.hover_z_m), int(round(trace.hover_time_s * 1e3))),
        family="monospace", size="x-small",
        color=hover_color, horizontalalignment="left", verticalalignment="bottom")
#     ax.text(
#         m_to_nm(trace.hover_z_m), vertical_segment_lower_limit - 10,
#         "%.1f nm" % (m_to_nm(trace.hover_z_m)), family="monospace", size="x-small", color=hover_color,
#         horizontalalignment="center", verticalalignment="top")
    ax.set_xlabel("Piezo position [nm]")
    ax.set_ylabel("Force [pN]")
    ax.set_ylim(vertical_segment_lower_limit - 100, vertical_segment_upper_limit + 300)

def createOnlyRefoldingFigure(trace_path, wlc_high_n, ignored_peaks_at_end, image):
    figure = _create_figure()
    refolding_ax = figure.add_subplot(111)
    with open(trace_path, "rt") as f:
        l = list(f)
        common.check_equal(len(l), 1)
    refolding_sample_curve(l[0], wlc_high_n, ignored_peaks_at_end, refolding_ax)
    canvas = matplotlib.backends.backend_ps.FigureCanvasPS(figure)
    canvas.print_figure(image)

def read_two_columns(file_object):
    a = []
    for line in file_object:
        posNm, timeMs = map(float, line.split())
        a.append((posNm, timeMs))
    x, y = reversed(zip(*a))
    return x, y

waiting_start_end_pattern = re.compile(
    r"WaitingStart=(?P<WaitingStart>[^ ]+) WaitingEnd=(?P<WaitingEnd>[^ ]+)")

def createZtFigure(zt_path, filename):
    figure = _create_figure()
    ax = figure.add_subplot(111)
    with open(zt_path, "rt") as f:
        lines = f.readlines()
        m = waiting_start_end_pattern.match(lines[0])
        waiting_start_s, waiting_end_s = m.groups()
        x, y = read_two_columns(lines[1:])
    ax.plot(x, y, linestyle="-", color="blue")
    ax.plot([0, max(x)], [0, 0], linestyle="dashed", color="black")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Piezo position (nm)")
#     ax.set_xlim(-max(x) / 50.0, max(x) * 1.1)
#     ax.set_ylim(min(y), max(y))
    ax.plot([waiting_start_s, waiting_start_s], [min(y), max(y)], linestyle="dotted", color="red")
    ax.plot([waiting_end_s, waiting_end_s], [min(y), max(y)], linestyle="dotted", color="red")
    ax.set_xlim(min(x), max(x))
    ax.set_ylim(min(y), max(y))
    canvas = matplotlib.backends.backend_ps.FigureCanvasPS(figure)
    canvas.print_figure(filename)

def main(params):
    """
    python refolding_trace.py --trace_path=/Users/daniel/work/curves/refolding/exported_traces/06111931.013.txt --wlc_high_pn=250 --ignored_peaks_at_end=0 --image_path=sampleRefoldingCurve.eps
    """
    parser = optparse.OptionParser()
    parser.add_option(
        "--trace_path",
        help="optionall trace file created using it.unibo.refolding.pico.stripchart.TraceExporter")
    parser.add_option("--zt_path", help="optional zt file created with it.unibo.refolding.pico.stripchart.TraceExporter")
    parser.add_option("--wlc_high_pn", help="until what force should wlc extend")
    parser.add_option("--image_path", help="path of eps file")
    parser.add_option("--zt_image_path", help="path of zt eps file")
    parser.add_option("--ignored_peaks_at_end", type=int, help="how many peaks to ignore at end")
    
    options, args = parser.parse_args(params)
    if args:
        parser.print_help()
        return -1
    print(options)

    pylab.rc("font", family="serif")
    pylab.rc("font", size=10)
    pylab.rc("axes", linewidth=0.5)
    pylab.rc('legend', fontsize='small')
    
    if options.trace_path:
        createOnlyRefoldingFigure(
            options.trace_path, float(options.wlc_high_pn) * 1e-12,
            int(options.ignored_peaks_at_end), options.image_path)
    if options.zt_path:
        createZtFigure(options.zt_path, options.zt_image_path)

if __name__ == '__main__':
    main(sys.argv[1:])
